"""
attention:
Q  K  V
"""
import torch
from torch import nn


class GxlOneConvMlp(nn.Module):
    """
    多层感知机，使用一维卷积进行线性变换，
    输入shape为:(b,seq,input_dim)
    输出shape: (b,seq,output_dim)
    示例：
    model = GxlOneConvMlp()
    input = torch.randn(10, 4, 1200)
    print(model(input).shape)
    > torch.Size([10, 4, 252])
    """

    def __init__(self, input_dim=1200, hidden_dim=1024, output_dim=252, dropout=0.0897):
        super(GxlOneConvMlp, self).__init__()
        self.conv1 = nn.Conv1d(in_channels=input_dim, out_channels=hidden_dim, kernel_size=1)
        self.conv2 = nn.Conv1d(in_channels=hidden_dim, out_channels=output_dim, kernel_size=1)
        self.ac_f = nn.ReLU()
        self.dropout = nn.Dropout(dropout)

    def forward(self, X: torch.Tensor):
        """
        X : (b,seq,emb)
        """
        X = self.conv2(self.ac_f(self.conv1(X.permute(0, 2, 1)))).permute(0, 2, 1)
        return self.dropout(X)


class GxlLinearMlp(nn.Module):
    """
    用线性模型实现的多层感知机
    输入shape为:(b,seq,input_dim)
    输出shape: (b,seq,output_dim)
    该实现模式本质上和一维卷积的实现模式相同，均是对input的最后一个维度做线性变换
    示例：
    model = GxlLinearMlp()
    input = torch.randn(10, 4, 1200)
    print(model(input).shape)
    > torch.Size([10, 4, 252])
    """

    def __init__(self, input_dim=252, hidden_dim=1024, output_dim=252, dropout=0.0897):
        super(GxlLinearMlp, self).__init__()
        self.linear1 = nn.Linear(input_dim, hidden_dim)
        self.linear2 = nn.Linear(hidden_dim, output_dim)
        self.ac_f = nn.ReLU()
        self.dropout = nn.Dropout(dropout)

    def forward(self, X):
        """
        X :(b,seq,emb)
        """
        X = self.linear2(self.ac_f(self.linear1(X)))
        return self.dropout(X)


class GxlMultiHeadAttention(nn.Module):
    """
    自定义多头注意力机制模块
    由注意力块和前馈多层感知机组成
    残差结构，输出为输入形状
    输入：
    key value是同一个量， query可以不同，若相同，则是self-attention,若不同，则是cross-attention
    mask: 决定哪些value是无用的，等到系数计算完毕后，将其系数置为0，（若采用softmax,则置为极小负值）
    query: (b, seq2, emb)
    key: (b, seq, emb)
    value: (b, seq1, emb)
    mask: (b, seq2, seq1)
    无需考虑batch_size,这只是批量计算的表征罢了。
    mask代表了对于指定的query，被查询的seq_list 中有哪几个单元是无效的
    输出： (b, seq2, emb)
    """

    def __init__(self, head_num=12, embedding_dim=252, representative_dim=1200, dropout=0.0897):
        super(GxlMultiHeadAttention, self).__init__()
        self.dim_per_head = representative_dim // head_num
        self.head_num = head_num
        self.q_convertor = nn.Linear(embedding_dim, representative_dim)
        self.k_convertor = nn.Linear(embedding_dim, representative_dim)
        self.v_convertor = nn.Linear(embedding_dim, representative_dim)
        self.mlp = GxlOneConvMlp(input_dim=representative_dim, hidden_dim=1024, output_dim=embedding_dim,
                                 dropout=dropout)
        self.dropout = nn.Dropout(p=dropout)
        self.layerNorm_f = nn.LayerNorm(normalized_shape=embedding_dim)  # 内部存放可学习的参数，也就是说一个标准化函数只能用在一个地方

    def forward(self, query, key, value, mask):
        """
        key value是同一个量， query可以不同，若相同，则是self-attention,若不同，则是cross-attention
        mask: 决定哪些value是无用的，等到系数计算完毕后，将其系数置为0，（若采用softmax,则置为极小负值）
        key: (b, seq1, emb)
        value: (b, seq1, emb)
        query: (b, seq2, emb)
        mask: (b, seq2, seq1)
        无需考虑batch_size,这只是批量计算的表征罢了。
        mask也就代表了对于指定的query，被查询的seq_list 中有哪几个单元是无效的，暂时不提供区分head的mask机制
        """
        residual = query
        batch_size = key.shape[0]
        seq2 = query.shape[1]

        # 得到q k v的代表人
        # representative_k:  (b,head_num,seq1,pre_small); representative_q:  (b,head_num,seq2,pre_small)
        representative_k = self.k_convertor(key) \
            .reshape(batch_size, -1, self.head_num, self.dim_per_head).permute(0, 2, 1, 3)
        representative_v = self.v_convertor(value) \
            .reshape(batch_size, -1, self.head_num, self.dim_per_head).permute(0, 2, 1, 3)
        representative_q = self.q_convertor(query) \
            .reshape(batch_size, -1, self.head_num, self.dim_per_head).permute(0, 2, 1, 3)
        #  (b,head_num,seq2,pre_small) * (b,head_num,pre_small, seq1) = weight : (b, h, seq2, seq1)

        # 计算得到权重值
        weight = torch.matmul(representative_q, representative_k.permute(0, 1, 3, 2))

        #  apply mask
        # mask (b,1,seq2, seq1) 方便广播，这里mask认为所有头都一样，暂时不提供区分head的mask机制，
        # 若有需要，请使用nn.MultiHeadAttention类，懒得费脑子实现了，目前我是暂时用不到
        if mask is not None:
            mask = torch.unsqueeze(mask, dim=1)
            weight.masked_fill(mask, -float('inf'))  # 利用广播机制，weight[mask]=.的形式不会自动广播
            weight = nn.functional.softmax(weight, dim=3)

        # weighted_res = (b, h ,seq2, pre_small) =(b, h, seq2, seq1) * (b, h, seq1, pre_small)
        weighted_res = torch.matmul(weight, representative_v)
        # weighted_res_no_head: (b ,seq2, pre)
        weighted_res_no_head = weighted_res.permute(0, 2, 1, 3).reshape(batch_size, seq2, -1)
        res = self.dropout(self.mlp(weighted_res_no_head))
        res = self.layerNorm_f(res + residual)
        return res


def get_padding_mask(value_data: torch.Tensor, pad_id, query_seq=None):
    """
    得到padding mask,

    当query_num为None时，代表self-attention
    当query_num为-1时， 代表无广播形式返回，返回数据为(batch,seq1)

    当value_data为时(b,)时，代表输入的是长度值，直接省去计算哪里为padding的那一步

    输入:
    value_data:(b, seq1)|(b,),
    pad_id: int,
    query_num: int
    输出： -》 (b, seq1)_bool -》（广播）(b, query_seq, seq1), padding_mask与查询无关
    """
    device = value_data.device
    if value_data.dim() == 2:  # 代表valid_data是no_embedding的value值：（batch,seq1),
        if query_seq is None:  # 代表self-attention
            query_seq = value_data.shape[1]
        elif query_seq == -1:
            return value_data == pad_id
        mask = (value_data == pad_id).to(device)
        mask = mask.unsqueeze(1).expand(-1, query_seq, -1)
        return mask
    elif value_data.dim() == 1:  # 代表输入的是长度:(batch,)
        max_value, _ = torch.max(value_data, dim=0)
        seq1 = max_value.item()
        temp = torch.arange(seq1).unsqueeze(0).expand(value_data.shape[0], -1).to(device)
        lens = value_data.unsqueeze(1).to(device)  # 0,1,2,3,4 |  3
        mask = temp >= lens  # (batch, seq1)
        if query_seq is None:  # 代表self-attention
            mask = mask.unsqueeze(1).expand(-1, seq1, -1)
            return mask.to(device)
        elif query_seq == -1:
            return mask.to(device)


def get_seq_mask(seq: int, batch_size: int):
    """
    得到seq mask,用于 masked self-attention ，屏蔽当前时间步之后的信息
    输入：seq(int) 为query 的句长，(因为是self-attention,所以也是kv的句长),batch_size
    输出：(seq,seq)下三角（不包括对角线）为false,上面的为true,->(广播)->(batch_size,seq,seq),seq_mask与batch无关
    """
    boundary = torch.arange(seq).unsqueeze(1)
    nums = torch.arange(seq).unsqueeze(0).expand(seq, -1)
    mask = nums >= boundary
    mask = mask.unsqueeze(0).expand(batch_size, -1, -1)
    return mask


if __name__ == "__main__":
    """"""
    # print(GxlMultiHeadAttention())
    # print(torch.ones(3, 3).bool())
    lens = torch.randint(0, 400, (10,))
    print(lens)
    maxvalue, _ = torch.max(lens, dim=0)
    print(maxvalue)
